# 计算每项数据的空气质量指数
class AQI:
    def __init__(self, values):
        self.pollutants = ["PM2.5", "PM10", "SO2", "NO2", "CO", "O3"]
        self.values = values
        # 空气质量分指数及对应的污染物项目脓限值，除了O3（臭氧）外，其他污染物项目浓度限值均为日均值
        self.table_day = {
            "PM2.5": [0, 35, 75, 115, 150, 250, 350, 500],
            "PM10": [0, 50, 150, 250, 350, 420, 500, 600],
            "SO2": [0, 50, 150, 475, 800, 1600, 2100, 2620],
            "NO2": [0, 40, 80, 180, 280, 565, 750, 940],
            "CO": [0, 2, 4, 14, 24, 36, 48, 60],
            "O3": [0, 160, 200, 300, 400, 800, 1000, 1200]  # 1小时浓度限值
        }
        # 空气质量分指数及对应的污染物项目浓限值，PM2.5和PM10没有小时平均浓限值
        self.table_hour = {
            "SO2": [0, 150, 500, 650, 800],
            "NO2": [0, 100, 200, 700, 1200, 2340, 3090, 3840],
            "CO": [0, 5, 10, 35, 60, 90, 120, 150],
            "O3": [0, 160, 200, 300, 400, 800, 1000, 1200]  # 1小时浓度限值
        }
        self.table_hour_eight = {
            "O3": [0, 100, 160, 215, 265, 800]  # 8小时浓度限值
        }
        self.IAQI = [0, 50, 100, 150, 200, 300, 400, 500]

    def get_IAQ(self, hour=False):
        each_IAQI = []
        for index in range(len(self.pollutants)):
            if hour:
                each_IAQI.append(self.hour_calculate(self.values[index], self.pollutants[index]))
            else:
                each_IAQI.append(self.day_calculate(self.values[index], self.pollutants[index]))

        compare = [0, 50, 100, 150, 200, 300]
        features = [("一级", "优", "绿色"), ("二级", "良", "黄色"), ("三级", "轻度污染", "橙色"), ("四级", "中度污染", "红色"),
                    ("五级", "重度污染", "紫色"), ("六级", "严重污染", "褐红色")]

        try:
            aqi = max(each_IAQI)
        except TypeError as e:
            print("PM2.5和PM10没有小时平均浓限值，无法计算分指数，返回其他污染物的分指数")
            return dict(zip(["SO2", "NO2", "CO", "O3"], each_IAQI[2:]))
        the_pollutant = ["PM2.5", "PM10", "SO2", "NO2", "CO", "O3"][each_IAQI.index(aqi)]
        for i in range(len(compare) - 1):
            if compare[i] <= aqi < compare[i + 1] + 1 and i == 0:
                return {
                    "AQI": aqi,
                    # "级别": features[i][0],
                    # "类别": features[i][1],
                    # "颜色": features[i][2],
                    "首要污染物": "无"
                }
            elif compare[i] <= aqi < compare[i + 1] + 1 and i == 1:
                return {
                    "AQI": aqi,
                    # "级别": features[i][0],
                    # "类别": features[i][1],
                    # "颜色": features[i][2],
                    "首要污染物": the_pollutant
                }
            elif compare[i] <= aqi < compare[i + 1] + 1 and i >= 2:
                return {
                    "AQI": aqi,
                    # "级别": features[i][0],
                    # "类别": features[i][1],
                    # "颜色": features[i][2],
                    "超标污染物": the_pollutant
                }
        return {
            "AQI": aqi,
            # "级别": features[5][0],
            # "类别": features[5][1],
            # "颜色": features[5][2],
            "超标污染物": the_pollutant
        }

    def day_calculate(self, C_p, pollutant):
        if C_p > self.table_day[pollutant][len(self.table_day[pollutant]) - 1]:
            return 500
        if C_p == 0:
            return 0
        for i in range(len(self.table_day[pollutant]) - 1):
            if self.table_day[pollutant][i] < C_p <= self.table_day[pollutant][i + 1]:
                BP_hi = self.table_day[pollutant][i + 1]
                BP_lo = self.table_day[pollutant][i]
                IAQI_hi = self.IAQI[i + 1]
                IAQI_lo = self.IAQI[i]
                IAQI_p = self.cal_linear(IAQI_hi, IAQI_lo, BP_hi, BP_lo, C_p)
                return IAQI_p

    def hour_calculate(self, C_p, pollutant):
        if pollutant == 'PM2.5' or pollutant == 'PM10':
            return None
        if pollutant == "SO2" and C_p > 800:
            print("二氧化硫小时平均脓限值超过800，需进行日平均浓度的分指数计算")
            return 6666
        if C_p > self.table_hour[pollutant][len(self.table_hour[pollutant]) - 1]:
            return 500
        if C_p == 0:
            return 0
        for i in range(len(self.table_hour[pollutant]) - 1):
            if self.table_hour[pollutant][i] < C_p <= self.table_hour[pollutant][i + 1]:
                BP_hi = self.table_hour[pollutant][i + 1]
                BP_lo = self.table_hour[pollutant][i]
                IAQI_hi = self.IAQI[i + 1]
                IAQI_lo = self.IAQI[i]
                IAQI_p = self.cal_linear(IAQI_hi, IAQI_lo, BP_hi, BP_lo, C_p)
                return IAQI_p

    def eight_hour(self, C_p, pollutant):
        if pollutant == "O3" and C_p > 800:
            print("臭氧8小时平均浓限值超过800，已按照1小时平均计算分指数")
            return self.hour_calculate(C_p, pollutant)
        if C_p > self.table_hour_eight[pollutant][len(self.table_hour_eight[pollutant]) - 1]:
            return 500
        if C_p == 0:
            return 0
        for i in range(len(self.table_hour_eight[pollutant]) - 1):
            if self.table_hour_eight[pollutant][i] < C_p <= self.table_hour_eight[pollutant][i + 1]:
                BP_hi = self.table_hour_eight[pollutant][i + 1]
                BP_lo = self.table_hour_eight[pollutant][i]
                IAQI_hi = self.IAQI[i + 1]
                IAQI_lo = self.IAQI[i]
                IAQI_p = self.cal_linear(IAQI_hi, IAQI_lo, BP_hi, BP_lo, C_p)
                return IAQI_p

    @staticmethod
    # 污染物空气质量分指数计算
    def cal_linear(IAQI_hi, IAQI_lo, BP_hi, BP_lo, C_p):
        '''
        :param IOQI_hi: 表中BP_hi对应的空气质量指数
        :param IAOI_lo: 表中BP_lo对应的空气质量指数
        :param BP_hi:   表中与C_p相近的污染物脓度限值的高位值
        :param BP_lo:   表中与C_p相近的污染物脓度限值的低位值
        :param C_p:     污染物项目的质量浓度值
        :return:    污染物项目的质量浓度
        '''
        # 线性缩放
        IAQI_p = (IAQI_hi - IAQI_lo) * (C_p - BP_lo) / (BP_hi - BP_lo) + IAQI_lo
        return int(new_round(IAQI_p, 0))
def new_round(_float, _len):
    '''
    四舍五入
    :param _float: 浮点数
    :param _len:    保留的小数位
    :return:
    '''
    if isinstance(_float, float):
        if str(_float)[::-1].find('.') <= _len:
            return (_float)
        if str(_float)[-1] == '5':
            return (round(float(str(_float)[:-1] + '6'), _len))
        else:
            return (round(_float, _len))
    else:
        return (round(_float, _len))

if __name__ == '__main__':
    h_aqi = AQI([24.57, 33.33, 9.97, 4.43, 0.35, 58.29]).get_IAQ(hour=True)
    print(h_aqi)
